# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import json
import time
import collections
import datetime
import numpy
import logging
from absl import app, flags
from coin.base.timestamp import get_timestamp
from coin.strategy.sync_rest_executor.util import *
from coin.strategy.sync_rest_executor.exec_scheduler import *
from coin.strategy.kr_hard_arb.util import ThreadRunner, execute_thread_runners
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str

flags.DEFINE_string('config_filename', '', '')

# if you need to change the variable name to consolidate, please do so.
flags.DEFINE_string('og_config_filename', '', '')

flags.DEFINE_bool('dry_run', True, '')

example_config = {
    'exchange':
        'Liqui',
    'targets': [{
        'base_currency': 'BTC',
        'quote_currency': 'ANT',  # target qty is of course, 0.
        'target_qty': 0,
        # fixed average lot size.
        # we have 290K ANT, so..probably divide in (48h/2m)
        # 48*60/2=1440
        # 290K/1440 ~= 201
        'average_lot_size': 201,
        'order_update_period_sec': 10,  # 'agg' or 'pass'
        'execution_mode': 'agg'
    }]
}


class SymbolExecutor(object):
  def __init__(self,
               exchange,
               *,
               base_currency,
               quote_currency,
               target_qty,
               average_lot_size,
               min_lot_size,
               order_update_period_sec,
               execution_mode,
               payup=0,
               max_payup_bps=None):
    self.product = generate_product_from_str(exchange=exchange,
                                             api_version=None,
                                             norm_product="%s-%s" % (base_currency, quote_currency))
    self.exec_scheduler = ExecScheduler(quote_currency=quote_currency,
                                        product=self.product,
                                        target_qty=target_qty,
                                        average_lot_size=average_lot_size,
                                        min_lot_size=min_lot_size)
    self.order_update_period = order_update_period_sec * 1e9
    # None, 'submit', 'cancel'
    self.execution_mode = execution_mode
    self.last_event = None
    self.last_submit_time = 0
    self.last_cancel_time = 0
    self.payup = payup
    self.max_payup_bps = max_payup_bps
    self.order_update_period_instance = None
    self.num_orders = None

  def get_order_update_period_instance(self):
    if self.order_update_period_instance is None:
      self.order_update_period_instance = numpy.random.lognormal(
          numpy.log(self.order_update_period), 0.5)
      logging.info("update period is %s", self.order_update_period_instance)
    return self.order_update_period_instance

  def reset_update_period_instance(self):
    self.order_update_period_instance = None

  def on_book_holder(self, book_holder):
    self.agg_price_tuple = book_holder.get_price_by_qty(self.exec_scheduler.average_lot_size)
    self.bbo_price_tuple = book_holder.get_best_price()

  def set_num_orders(self, order_base):
    lst = order_base.query_list_orders(self.product)
    self.num_orders = len(lst.msg.each_order)

  def poll_num_orders(self, order_base):
    execute_thread_runners([ThreadRunner(self.set_num_orders, order_base)], "poll num of orders")

  def trade(self, order_base, balance):
    self.exec_scheduler.update_balance(balance)
    order_side, order_qty = self.exec_scheduler.get_order_side_and_qty()
    timestamp = get_timestamp()
    logging.info("last event: %s, num orders of %s: %s",
                 self.last_event,
                 self.product,
                 self.num_orders)
    if order_qty is not None and order_qty > 0:
      if (self.last_event is None or (self.last_event == 'cancel' and self.num_orders == 0)):
        if (timestamp - self.last_cancel_time > 5e8 and self.agg_price_tuple is not None
            and self.bbo_price_tuple is not None):
          order_price = get_order_price(execution_mode=self.execution_mode,
                                        order_side=order_side,
                                        agg_price_tuple=self.agg_price_tuple,
                                        bbo_price_tuple=self.bbo_price_tuple,
                                        payup=self.payup,
                                        max_payup_bps=self.max_payup_bps)
          if order_price is None:
            logging.info("payup limit exceeds, %s, %s, %s",
                         self.product,
                         self.agg_price_tuple,
                         self.bbo_price_tuple)
            return
          assert order_price is not None and order_price > 0, (self.agg_price_tuple,
                                                               self.bbo_price_tuple)
          self.num_orders = None
          self.last_submit_time = timestamp
          self.last_event = 'submit'
          params = [self.product, order_side_to_sync_side(order_side), order_price, order_qty]
          if flags.FLAGS.dry_run:
            logging.info("should submit %s side %s at %.8f of %s, but it's dry run", *params)
          elif self.execution_mode == 'agg' and hasattr(order_base, "query_submit_ioc"):
            logging.info("submit ioc %s side %s at %.8f of %s", *params)
            execute_thread_runners([ThreadRunner(order_base.query_submit_ioc, *params)],
                                   "submit ioc order")
          elif self.execution_mode in ['pass', 'pass_deep']:
            logging.info("submit %s side %s at %.8f of %s", *params)
            execute_thread_runners([ThreadRunner(order_base.query_submit, *params)], "submit order")
      elif self.last_event == 'submit' or (self.last_event == 'cancel' and self.num_orders > 0
                                           and timestamp - self.last_cancel_time > 5e8):
        if self.last_event == 'submit':
          update_period = self.get_order_update_period_instance()
          if timestamp - self.last_submit_time > update_period:
            self.reset_update_period_instance()
          else:
            return
        self.poll_num_orders(order_base)
        if self.execution_mode == 'agg':
          # ioc style
          self.last_cancel_time = timestamp
          self.last_event = 'cancel'
        else:
          self.last_cancel_time = timestamp
          self.last_event = 'cancel'
          if flags.FLAGS.dry_run:
            logging.info("should cancel all for %s, but it's dry run", self.product)
          else:
            logging.info("cancel all for %s", self.product)
            execute_thread_runners([ThreadRunner(order_base.query_cancel_product, self.product)],
                                   "cancel order")
      else:
        raise ValueError(self.last_event)
    return order_side, order_qty


class VolumeExecutor(object):
  def __init__(self, execution_spec):
    self._symbol_executors = collections.OrderedDict()
    self._balance_container = BalanceContainer()
    for symbol_config in execution_spec['targets']:
      symbol_executor = SymbolExecutor(execution_spec['exchange'], **symbol_config)
      if symbol_executor.product not in self._symbol_executors:
        self._symbol_executors[symbol_executor.product] = []
      self._symbol_executors[symbol_executor.product].append(symbol_executor)

  def get_products(self):
    return self._symbol_executors.keys()

  def on_book_holder(self, product, book_holder):
    if product in self._symbol_executors:
      for symbol_exec in self._symbol_executors[product]:
        symbol_exec.on_book_holder(book_holder)

  def on_balance(self, balance_proto):
    self._balance_container.update_balance(get_timestamp(), balance_proto)

  def trade_product(self, product, order_base):
    if product in self._symbol_executors:
      balance = self._balance_container.get_balance(product.base.currency)
      if balance is not None:
        for symbol_exec in self._symbol_executors[product]:
          symbol_exec.trade(order_base, balance)


InfoBase = collections.namedtuple('InfoBase', ['product', 'order_base', 'feed_base'])


class SyncRestExecutor(VolumeExecutor):
  def __init__(self, execution_spec):
    VolumeExecutor.__init__(self, execution_spec)
    self.bases = []
    self.last_order_poll = 0
    for product in self.get_products():
      order_base, feed_base = get_order_feed_base("Spot", execution_spec['exchange'], product,
                                                  flags.FLAGS.og_config_filename)
      self.bases.append(InfoBase(product, order_base, feed_base))

  def poll(self):
    timestamp = get_timestamp()
    for info_base in self.bases:
      execute_thread_runners([ThreadRunner(info_base.feed_base.query_book)], "query book")

    if self.last_order_poll + 10e9 < timestamp:
      for info_base in self.bases:
        execute_thread_runners([ThreadRunner(info_base.order_base.query_balance)], "query balance")
      self._balance_container.print_balance()
      self.last_order_poll = timestamp

  def trade(self):
    for info_base in self.bases:
      product = info_base.product
      book_holder = info_base.feed_base
      self.on_book_holder(product, book_holder)
      bal_cached = info_base.order_base.get_balance_cached()
      if bal_cached is None:
        continue
      self.on_balance(bal_cached.msg)

    for info_base in self.bases:
      self.trade_product(info_base.product, info_base.order_base)


def main(_):
  with open(flags.FLAGS.config_filename, 'r') as conf_file:
    conf = json.load(conf_file)
    sync_res_exec = SyncRestExecutor(conf)
    while (True):
      try:
        time.sleep(2)
        sync_res_exec.poll()
        logging.info(datetime.datetime.now())
        sync_res_exec.trade()
      except Exception:
        import traceback
        traceback.print_exc()
  return 0


if __name__ == '__main__':
  app.run(main)
